from mindspore import numpy as mnp, Tensor
from mindspore import nn, ops, dtype as mstype

from transformers_network.multi_head_attention import MultiHeadAttention
from transformers_network.utils import get_attn_pad_mask


class PositionalEncoding(nn.Cell):
    """位置编码
    用于处理序列数据（如文本）以便能让模型识别单词在序列中的相对位置
    """

    def __init__(self, d_model, dropout_p=0.1, max_len=100):
        """
        位置编码是一种通过对输入序列的位置信息进行编码的方式，使得模型能够更好地捕获序列中的位置关系。
        位置编码的公式为：
            PE(pos,2i) = sin(pos/10000^(2i/d_model))
            PE(pos,2i+1) = cos(pos/10000^(2i/d_model))
        其中pos表示序列的位置，2i和2i+1表示位置的偶数和奇数分量，d_model表示模型的维度。
        位置编码的作用是使得模型能够学习到输入序列的位置信息，从而能够捕获到不同位置之间的关系。
        """
        super().__init__()
        self.dropout = nn.Dropout(p=dropout_p)

        self.pe = ops.Zeros()((max_len, d_model), mstype.float32)  # 初始化位置编码

        pos = mnp.arange(0, max_len, dtype=mstype.float32).view((-1, 1))  # 生成位置序列[0, 1, 2, 3, ...]
        angle = ops.pow(10000.0, mnp.arange(0, d_model, 2, dtype=mstype.float32) / d_model)

        # 使用正弦和余弦函数为每个位置生成 唯一的 编码。该编码可以帮助模型理解元素在序列中的位置信息。
        self.pe[:, 0::2] = ops.sin(pos / angle)
        self.pe[:, 1::2] = ops.cos(pos / angle)

    def construct(self, x):
        """
        :param x: 经过了embedding后的输入序列
        :return:  x加上位置编码信息后，经过dropout后输出
        """
        batch_size = x.shape[0]

        pe = self.pe.expand_dims(0)
        pe = ops.broadcast_to(pe, (batch_size, -1, -1))

        x = x + pe[:, :x.shape[1], :]
        return self.dropout(x)


x = ops.Zeros()((1, 2, 4), mstype.float32)
pe = PositionalEncoding(4)
pe_x = pe(x)
print(pe_x.shape)
print(pe_x)


class PoswiseFeedForward(nn.Cell):
    def __init__(self, d_ff, d_model, dropout_p=0.):
        """
        基于位置的前馈神经网络被用来对输入中的每个位置进行非线性变换。它由两个线性层组成，层与层之间需要经过ReLU激活函数。
         FFN(𝑥)=ReLU(𝑥𝑊1 + 𝑏1)𝑊2 + 𝑏2
        相比固定的ReLU函数，基于位置的前馈神经网络可以处理更加复杂的关系，并且由于前馈网络是基于位置的，可以捕获到不同位置的信息，并为每个位置提供不同的转换。
            `d_ff`：表示前馈神经网络的隐藏层大小。
            `d_model`：表示输入的特征维度。
            `dropout_p`：表示dropout的概率。
        """
        super().__init__()
        self.linear1 = nn.Dense(d_model, d_ff)
        self.linear2 = nn.Dense(d_ff, d_model)
        self.dropout = nn.Dropout(p=dropout_p)
        self.relu = nn.ReLU()

    def construct(self, x):
        """前馈神经网络
        x: [batch_size, seq_len, d_model]
        """
        x = self.linear1(x)
        x = self.relu(x)
        x = self.dropout(x)
        output = self.linear2(x)
        return output


# x = ops.ones((1, 2, 4), mstype.float32)
# ffn = PoswiseFeedForward(16, 4)
# print(ffn(x).shape)

class AddNorm(nn.Cell):
    def __init__(self, d_model, dropout_p=0.):
        """
            Add&Norm(𝑥) = LayerNorm(𝑥 + Sublayer(𝑥))
            Add：残差连接，帮助缓解网络退化问题，注意需要满足𝑥与SubLayer(𝑥)的形状一致；
            Norm：LayerNorm，层归一化，帮助模型更快地进行收敛；
        """
        super().__init__()
        self.layer_norm = nn.LayerNorm((d_model,), epsilon=1e-5)
        self.dropout = nn.Dropout(p=dropout_p)

    def construct(self, x, residual):
        """残差连接+层归一化
        x: [batch_size, seq_len, d_model]
        residual: [batch_size, seq_len, d_model]
        """
        return self.layer_norm(self.dropout(x) + residual)


# x = ops.ones((1, 2, 4), mstype.float32)
# residual = ops.ones((1, 2, 4), mstype.float32)
# add_norm = AddNorm(4)
# print(add_norm(x, residual).shape)


class EncoderLayer(nn.Cell):
    def __init__(self, d_model, n_heads, d_ff, dropout_p=0.):
        super().__init__()

        self.enc_self_attn = MultiHeadAttention(d_model, n_heads, dropout_p)
        self.pos_ffn = PoswiseFeedForward(d_ff, d_model, dropout_p)
        self.add_norm1 = AddNorm(d_model, dropout_p)
        self.add_norm2 = AddNorm(d_model, dropout_p)

    def construct(self, enc_inputs, enc_self_attn_mask):
        """
        enc_inputs: [batch_size, src_len, d_model], 输入序列
        enc_self_attn_mask: [batch_size, src_len, src_len], 自注意力掩码
        """
        residual = enc_inputs
        # 自注意力
        enc_outputs, attn = self.enc_self_attn(enc_inputs, enc_inputs, enc_inputs, enc_self_attn_mask)
        enc_outputs = self.add_norm1(enc_outputs, residual)

        residual = enc_outputs
        # 位置
        enc_outputs = self.pos_ffn(enc_outputs)
        enc_outputs = self.add_norm2(enc_outputs, residual)

        return enc_outputs, attn


# x = ops.ones((1, 2, 8), mstype.float32)
# mask = Tensor([False]).broadcast_to((1, 2, 2))
# encoder_layer = EncoderLayer(8, 4, 16)
# output, attn = encoder_layer(x, mask)
# print(output.shape, attn.shape)


class Encoder(nn.Cell):
    def __init__(self, src_vocab_size, d_model, n_heads, d_ff, n_layers, dropout_p=0.):
        """
            Encoder模块，用于处理输入序列，包括词嵌入、位置编码、多头注意力、前馈神经网络、残差连接和层归一化。
            `src_vocab_size`：表示输入序列的词表大小。
            `d_model`：表示输入的特征维度。
            `n_heads`：表示多头注意力的头数。
            `d_ff`：表示位置前馈神经网络FFN的隐藏层大小。
            `n_layers`：表示编码器的层数。
            `dropout_p`：表示dropout的概率。
        """
        super().__init__()
        self.src_emb = nn.Embedding(src_vocab_size, d_model)
        self.pos_emb = PositionalEncoding(d_model, dropout_p)
        self.layers = nn.CellList([EncoderLayer(d_model, n_heads, d_ff, dropout_p) for _ in range(n_layers)])
        self.scaling_factor = ops.Sqrt()(Tensor(d_model, mstype.float32))

    def construct(self, enc_inputs, src_pad_idx):
        """
        enc_inputs : [batch_size, src_len], 输入序列
        src_pad_idx : [batch_size, src_len], 输入序列的padding mask, 值为True的位置为pad(只为了保证统一的输入序列长度,而实际无意义的填充值)
        """

        # 词嵌入
        enc_outputs = self.src_emb(enc_inputs.astype(mstype.int32))
        # 获取位置编码信息,
        enc_outputs = self.pos_emb(enc_outputs * self.scaling_factor)
        # 自注意力掩码
        enc_self_attn_mask = get_attn_pad_mask(enc_inputs, enc_inputs, src_pad_idx)

        enc_self_attns = []
        for layer in self.layers:
            enc_outputs, enc_self_attn = layer(enc_outputs, enc_self_attn_mask)
            enc_self_attns.append(enc_self_attn)
        return enc_outputs, enc_self_attns
